import numpy as np
import torch
import pickle
import os
import re
from tqdm import tqdm
from datetime import datetime


def synchronize_videoandradar_folders(allvideofolder='./VIDEO/', 
                                      allradardatafolder='./RADAR/',
                                      videodt_format = "%Y%m%d%H%M",
                                      radardt_format = "%Y%m%d%H%M%S"):
    """
    function: Put video files and radar files with the same capture time into the same folder
    args: allvideofolder: the root path of video data
          allradardatafolder: the root path of radar data
          videodt_format: the tine format of video
          radardt_format: the time format of radar
        
    return: put two into one filefolder

    """
    # List video and radar folders
    allvideofolder_list = os.listdir(allvideofolder)
    allvideofolder_list = sorted(allvideofolder_list, key=lambda x: int(re.findall(r'\d+', x)[0]))

    allradardatafolder_list = os.listdir(allradardatafolder)
    allradardatafolder_list = sorted(allradardatafolder_list, key=lambda x: int(re.findall(r'\d+', x)[0]))

    # Find matching folders
    same_folder = []
    for i in range(len(allvideofolder_list)):
        video_strat_timedt = datetime.strptime(str(allvideofolder_list[i]), videodt_format)
        video_strat_timestamp = video_strat_timedt.timestamp()
        
        for j in range(len(allradardatafolder_list)):
            radar_strat_timedt = datetime.strptime(str(allradardatafolder_list[j]), radardt_format)
            radar_strat_timestamp = radar_strat_timedt.timestamp()

            if abs(video_strat_timestamp - radar_strat_timestamp) <= 100:
                same_folder.append([allvideofolder_list[i], allradardatafolder_list[j]])

    newpath = './SYNCHRONIZATION'

    for i in range(len(same_folder)):
        videofolderpath = os.path.join(allvideofolder, same_folder[i][0])
        radarfolderpath = os.path.join(allradardatafolder, same_folder[i][1])
        
        bothsensorspath = os.path.join(newpath, same_folder[i][0])
        newvideofolderpath = os.path.join(newpath, same_folder[i][0], 'VIDEO')
        newradarfolderpath = os.path.join(newpath, same_folder[i][0], 'RADAR')
        
        if not os.path.exists(bothsensorspath):
            os.mkdir(bothsensorspath)
            
            if not os.path.exists(newvideofolderpath):
                os.mkdir(newvideofolderpath)
                shutil.move(videofolderpath, newvideofolderpath)
            
            if not os.path.exists(newradarfolderpath):
                os.mkdir(newradarfolderpath)
                shutil.move(radarfolderpath, newradarfolderpath)

# Example usage:
# synchronize_folders('./VIDEO/', './RADAR/')

